#!/usr/bin/env python3

import argparse
import xarray as xr
from multiprocessing import Pool
from itertools import zip_longest, repeat
import metpy.calc as calc
from metpy.units import units
import numpy as np
import pandas as pd
import pendulum
from scipy.interpolate import interp1d
from glob import glob
import os
import sys
sys.path.append(f'{os.path.dirname(os.path.realpath(__file__))}/../../utils')
from utils import cli, parse_time, parse_config, create_kdtree, find_bbox, create_bilinear_wgts, apply_bilinear, has_key

default_zlev = np.linspace(100, 6000, 30) * units.m

def sim_profiler_ecwf(file_root, sid, lon, lat, zlev, start_time, end_time):
	cli.notice(f'Simulate PROFILER at {sid}@({lon},{lat}).')
	files = file_root + '/W_NAFP_C_ECMF_*'
	used_files = []
	# Assume the file name contains date time.
	times = []
	for file in sorted(glob(files)):
		tmp = os.path.basename(file)
		time = pendulum.datetime(year=int(tmp[14:18]), month=int(tmp[42:44]), day=int(tmp[44:46]), hour=int(tmp[46:48]))
		# f = xr.open_dataset(file, engine='cfgrib', backend_kwargs={
		# 	'filter_by_keys': {
		# 		'typeOfLevel': 'surface',
		# 		'edition': 1
		# 	},
		# 	'indexpath': ''
		# })
		# time = pendulum.parse(str(f.variables['time'].values))
		if start_time <= time <= end_time:
			times.append(time)
			used_files.append(file)
	if len(used_files) == 0:
		cli.error(f'Cannot find valid ECWF files!')
	f_sfc = xr.open_mfdataset(used_files, engine='cfgrib', backend_kwargs={
		'filter_by_keys': {
			'typeOfLevel': 'surface',
			'edition': 1
		},
		'indexpath': ''
	}, concat_dim='valid_time', combine='nested')
	f_plev = xr.open_mfdataset(used_files, engine='cfgrib', backend_kwargs={
		'filter_by_keys': {
			'typeOfLevel': 'isobaricInhPa',
			'edition': 1
		},
		'indexpath': ''
	}, concat_dim='valid_time', combine='nested')

	if type(zlev) != list: zlev = default_zlev

	t2   = f_sfc.t2m.interp(longitude=lon, latitude=lat, method='linear') * units.K
	d2   = f_sfc.d2m.interp(longitude=lon, latitude=lat, method='linear') * units.K
	ps   = f_sfc.sp .interp(longitude=lon, latitude=lat, method='linear') * units.Pa
	slp  = f_sfc.msl.interp(longitude=lon, latitude=lat, method='linear') * units.Pa
	sh2  = calc.specific_humidity_from_dewpoint(ps, d2)
	mr2  = calc.mixing_ratio_from_specific_humidity(sh2)
	# tv   = calc.virtual_temperature_from_dewpoint(ps, t2, d2)
	tv   = calc.virtual_temperature(t2, mr2)
	gam  = 0.0065
	rd   = 287.05
	g    = 9.80665
	zs   = np.array([float(float(tv) / gam * ((slp / ps)**(gam * rd / g) - 1))]) * units.m

	# Height levels
	# First, do horizontal interpolation.
	h    = f_plev.gh.interp(longitude=lon, latitude=lat, method='linear').values / g * units.m
	utmp = f_plev.u .interp(longitude=lon, latitude=lat, method='linear').values * units.mps
	vtmp = f_plev.v .interp(longitude=lon, latitude=lat, method='linear').values * units.mps
	# Second, do vertical interpolation.
	u = []; v = []
	for i in range(h.shape[0]): # time dimension
		u.append(interp1d(h[i,:], utmp[i,:], kind='linear', bounds_error=False, copy=False)(zlev))
		v.append(interp1d(h[i,:], vtmp[i,:], kind='linear', bounds_error=False, copy=False)(zlev))
	u = np.array(u) * units.mps
	v = np.array(v) * units.mps
	# Third, calculate other quantities.
	ws = calc.wind_speed(u, v)
	wd = calc.wind_direction(u, v)
	zlev_data = []
	for i in range(len(times)):
		for j in range(len(zlev)):
			if not np.isnan(u[i,j]):
				zlev_data.append({
					'sid' : sid,
					'lon' : lon,
					'lat' : lat,
					'time': times[i],
					'zs'  : zs[i]  .magnitude,
					'h'   : zlev[j].magnitude,
					'u'   : u [i,j].magnitude,
					'v'   : v [i,j].magnitude,
					'ws'  : ws[i,j].magnitude,
					'wd'  : wd[i,j].magnitude,
				})
	return pd.DataFrame(zlev_data)

def sim_profiler_wrf(file_root, sid, lon, lat, zlev, start_time, end_time):
	cli.notice(f'Simulate PROFILER at {sid}@({lon},{lat}).')
	files = file_root + '/wrfout*'
	used_files = []
	# Assume the file name contains date time.
	times = []
	for file in sorted(glob(files)):
		f = xr.open_dataset(file, engine='netcdf4')
		time = pendulum.from_format(str(f.Times[0].data)[2:-1], 'YYYY-MM-DD_HH:mm:SS')
		if start_time <= time <= end_time:
			times.append(time)
			used_files.append(file)
	if len(used_files) == 0: cli.error('Cannot find wrfout files within time range!')
	f = xr.open_mfdataset(used_files, concat_dim='Time', combine='nested')

	if type(zlev) != list: zlev = default_zlev

	xlon = f.XLONG[0,:,:]
	xlat = f.XLAT [0,:,:]
	xlon_u = f.XLONG_U[0,:,:]
	xlat_u = f.XLAT_U [0,:,:]
	xlon_v = f.XLONG_V[0,:,:]
	xlat_v = f.XLAT_V [0,:,:]
	kdtree = create_kdtree(xlon, xlat)
	kdtree_u = create_kdtree(xlon_u, xlat_u)
	kdtree_v = create_kdtree(xlon_v, xlat_v)
	bbox = find_bbox(kdtree, lon, lat)
	bbox_u = find_bbox(kdtree_u, lon, lat)
	bbox_v = find_bbox(kdtree_v, lon, lat)
	wgts_h = create_bilinear_wgts(xlon, xlat, bbox, lon, lat)
	wgts_h_u = create_bilinear_wgts(xlon_u, xlat_u, bbox_u, lon, lat)
	wgts_h_v = create_bilinear_wgts(xlon_v, xlat_v, bbox_v, lon, lat)

	zs = apply_bilinear(wgts_h, f.HGT, dims=('south_north','west_east')) * units.m

	# Height levels
	# First, do horizontal interpolation.
	ph   = 0.5 * (f.PH [:,0:f.PH .shape[1]-1,:,:] + f.PH [:,1:f.PH .shape[1],:,:]) / 9.8
	phb  = 0.5 * (f.PHB[:,0:f.PHB.shape[1]-1,:,:] + f.PHB[:,1:f.PHB.shape[1],:,:]) / 9.8
	h    = apply_bilinear(wgts_h, ph + phb, dims=('south_north', 'west_east'))
	utmp = apply_bilinear(wgts_h_u, f.U, dims=('south_north', 'west_east_stag'))
	vtmp = apply_bilinear(wgts_h_v, f.V, dims=('south_north_stag', 'west_east'))
	# Second, do vertical interpolation.
	u = []; v = []
	for i in range(h.shape[0]): # time dimension
		# FIXME: Why pint complains "ValueError: Cannot operate with Quantity and Quantity of different registries." when zs[0] + zlev?
		u.append(interp1d(h[i,:], utmp[i,:], kind='linear', bounds_error=False, copy=False)(zs[0].magnitude + zlev.magnitude))
		v.append(interp1d(h[i,:], vtmp[i,:], kind='linear', bounds_error=False, copy=False)(zs[0].magnitude + zlev.magnitude))
	u = np.array(u) * units.mps
	v = np.array(v) * units.mps
	# Third, calculate other quantities.
	ws = calc.wind_speed(u, v)
	wd = calc.wind_direction(u, v)
	zlev_data = []
	for i in range(len(times)):
		for j in range(len(zlev)):
			if not np.isnan(u[i,j]):
				zlev_data.append({
					'sid' : sid,
					'lon' : lon,
					'lat' : lat,
					'time': times[i],
					'zs'  : zs[i]  .magnitude,
					'h'   : zlev[j].magnitude,
					'u'   : u [i,j].magnitude,
					'v'   : v [i,j].magnitude,
					'ws'  : ws[i,j].magnitude,
					'wd'  : wd[i,j].magnitude,
				})
	return pd.DataFrame(zlev_data)

def add_obs_err(df, hlev_obs_err, obs_err_u, obs_err_v):
	tmp = pd.DataFrame({
		'h': hlev_obs_err,
		'u': obs_err_u,
		'v': obs_err_v
	}).set_index('h').sort_index()
	obs_err = pd.DataFrame(index=df['h'].unique()).sort_index()
	for key in tmp:
		obs_err[key] = np.interp(obs_err.index, tmp.index, tmp[key])
	for i in range(len(df)):
		h = df.loc[i,'h']
		df.loc[i,'u'] += np.random.normal(loc=0, scale=obs_err.loc[h,'u'], size=1)[0]
		df.loc[i,'v'] += np.random.normal(loc=0, scale=obs_err.loc[h,'v'], size=1)[0]
	df['ws'] = calc.wind_speed(list(df['u']) * units.mps, list(df['v']) * units.mps)
	df['wd'] = calc.wind_direction(list(df['u']) * units.mps, list(df['v']) * units.mps)
	return df

def write_profiler_littler(file_path, df):
	f = open(file_path, 'w')
	itime = 0
	for time in df['time'].unique():
		for sid in df.loc[df['time']==time]['sid'].unique():
			row = df.loc[df['time']==time].loc[df['sid']==sid].iloc[0]
			zlev = list(df.loc[df['time']==time].loc[df['sid']==sid].iterrows())
			# Header
			f.write('{:20.5f}' .format(row['lat']))                           # latitude
			f.write('{:20.5f}' .format(row['lon']))                           # longitude
			f.write('{:40}'    .format(row['sid']))                           # id
			f.write('{:40}'    .format(row['sid']))                           # name
			f.write('{:40}'    .format('FM-132 PROFLR'))                      # platform
			f.write('{:40}'    .format('OSSE'))                               # source
			f.write('{:>20.5f}'.format(row['zs']))                            # elevation
			f.write('{:>10d}'  .format(len(zlev)))                            # num_vld_fld
			f.write('{:>10d}'  .format(0))                                    # num_error
			f.write('{:>10d}'  .format(0))                                    # num_warning
			f.write('{:>10d}'  .format(itime))                                # seq_num
			f.write('{:>10d}'  .format(0))                                    # num_dups
			f.write('{:>10}'   .format('F'))                                  # is_sound
			f.write('{:>10}'   .format('F'))                                  # bogus
			f.write('{:>10}'   .format('F'))                                  # discard
			f.write('{:>10d}'  .format(-888888))                              # obs_time
			f.write('{:>10d}'  .format(-888888))                              # julian_day
			f.write('{:>20}'   .format(row['time'].strftime('%Y%m%d%H%M%S'))) # date time
			f.write('{:>13.5f}'.format(-888888.0))                            # slp
			f.write('{:>7d}'   .format(0))                                    # slp qc
			f.write('{:>13.5f}'.format(-888888.0))                            # ref_pres
			f.write('{:>7d}'   .format(0))                                    # ref_pres qc
			f.write('{:>13.5f}'.format(-888888.0))                            # ground_t
			f.write('{:>7d}'   .format(0))                                    # ground_t qc
			f.write('{:>13.5f}'.format(-888888.0))                            # sst
			f.write('{:>7d}'   .format(0))                                    # sst qc
			f.write('{:>13.5f}'.format(-888888.0))                            # psfc
			f.write('{:>7d}'   .format(0))                                    # psfc qc
			f.write('{:>13.5f}'.format(-888888.0))                            # precip
			f.write('{:>7d}'   .format(0))                                    # precip qc
			f.write('{:>13.5f}'.format(-888888.0))                            # t_max
			f.write('{:>7d}'   .format(0))                                    # t_max qc
			f.write('{:>13.5f}'.format(-888888.0))                            # t_min
			f.write('{:>7d}'   .format(0))                                    # t_min qc
			f.write('{:>13.5f}'.format(-888888.0))                            # t_max night
			f.write('{:>7d}'   .format(0))                                    # t_max night qc
			f.write('{:>13.5f}'.format(-888888.0))                            # p_tend03
			f.write('{:>7d}'   .format(0))                                    # p_tend03 qc
			f.write('{:>13.5f}'.format(-888888.0))                            # p_tend24
			f.write('{:>7d}'   .format(0))                                    # p_tend24 qc
			f.write('{:>13.5f}'.format(-888888.0))                            # cloud_cover
			f.write('{:>7d}'   .format(0))                                    # cloud_cover qc
			f.write('{:>13.5f}'.format(-888888.0))                            # ceiling
			f.write('{:>7d}'   .format(0))                                    # ceiling qc
			f.write('\n')
			ilev = 0
			for idx, row in zlev:
				# Record
				f.write('{:13.5f}'.format(-888888.0))                          # pressure (Pa)
				f.write('{:7d}'   .format(0))                                  # pressure qc
				f.write('{:13.5f}'.format(row['h']))                           # height
				f.write('{:7d}'   .format(0))                                  # height qc
				f.write('{:13.5f}'.format(-888888.0))                          # temperature (K)
				f.write('{:7d}'   .format(0))                                  # temperature qc
				f.write('{:13.5f}'.format(-888888.0))                          # dewpoint (K)
				f.write('{:7d}'   .format(0))                                  # dewpoint qc
				f.write('{:13.5f}'.format(row['ws']))                          # wind speed (m s-1)
				f.write('{:7d}'   .format(0))                                  # wind speed qc
				f.write('{:13.5f}'.format(row['wd']))                          # wind direction (deg)
				f.write('{:7d}'   .format(0))                                  # wind direction qc
				f.write('{:13.5f}'.format(row['u']))                           # wind u component (m s-1)
				f.write('{:7d}'   .format(0))                                  # wind u component qc
				f.write('{:13.5f}'.format(row['v']))                           # wind v component (m s-1)
				f.write('{:7d}'   .format(0))                                  # wind v component qc
				f.write('{:13.5f}'.format(-888888.0))                          # relative humidity (%)
				f.write('{:7d}'   .format(0))                                  # relative humidity qc
				f.write('{:13.5f}'.format(-888888.0))                          # thickness (m)
				f.write('{:7d}'   .format(0))                                  # thickness qc
				f.write('\n')
				ilev += 1
			# End
			f.write('{:13.5f}'.format(-777777.0))
			f.write('{:7d}'   .format(0))
			f.write('{:13.5f}'.format(-777777.0))
			f.write('{:7d}'   .format(0))
			f.write('{:13.5f}'.format(-888888.0))
			f.write('{:7d}'   .format(0))
			f.write('{:13.5f}'.format(-888888.0))
			f.write('{:7d}'   .format(0))
			f.write('{:13.5f}'.format(-888888.0))
			f.write('{:7d}'   .format(0))
			f.write('{:13.5f}'.format(-888888.0))
			f.write('{:7d}'   .format(0))
			f.write('{:13.5f}'.format(-888888.0))
			f.write('{:7d}'   .format(0))
			f.write('{:13.5f}'.format(-888888.0))
			f.write('{:7d}'   .format(0))
			f.write('{:13.5f}'.format(-888888.0))
			f.write('{:7d}'   .format(0))
			f.write('{:13.5f}'.format(-888888.0))
			f.write('{:7d}'   .format(0))
			f.write('\n')
			f.write('{:7d}{:7d}{:7d}\n'.format(ilev, 0, 0))
			itime += 1
	f.close()

def sim_profiler(nature_run_root, nature_run_type, sid, lon, lat, zlev, start_time, end_time, hlev_obs_err, obs_err_u, obs_err_v, ntasks, output_littler):
	# Check directory existence.
	if not os.path.isdir(os.path.dirname(os.path.abspath(output_littler))):
		cli.notice(f'Create directory {os.path.dirname(output_littler)}.')
		os.makedirs(os.path.dirname(os.path.abspath(output_littler)))

	# Do jobs in parallel.
	n = len(sid)
	pool = Pool(processes=ntasks)
	if nature_run_type == 'wrf':
		sim_profiler_call = sim_profiler_wrf
	elif nature_run_type == 'ecwf':
		sim_profiler_call = sim_profiler_ecwf
	res = pool.starmap(sim_profiler_call,
		zip_longest(
			repeat(nature_run_root, n),
			sid, lon, lat,
			repeat(zlev, n),
			repeat(start_time, n),
			repeat(end_time, n)
		)
	)

	# Concatenate results.
	df = pd.concat(res, ignore_index=True)

	# Add observation errors.
	df = add_obs_err(df, hlev_obs_err, obs_err_u, obs_err_v)

	# Write LITTLE_R file.
	write_profiler_littler(output_littler, df)

if __name__ == '__main__':
	parser = argparse.ArgumentParser('Simulate wind profiler observation.')
	parser.add_argument('-i', '--nature-run-root')
	parser.add_argument('-t', '--nature-run-type', choices=('ecwf', 'wrf'), required=True)
	parser.add_argument('-o', '--output-littler', required=True)
	parser.add_argument('--sid', nargs='+')
	parser.add_argument('--lon', nargs='+', type=float)
	parser.add_argument('--lat', nargs='+', type=float)
	parser.add_argument('--zlev', help='Height levels in meter', nargs='+', type=float)
	parser.add_argument('--start-time', required=True, type=parse_time)
	parser.add_argument('--end-time', required=True, type=parse_time)
	parser.add_argument('--hlev-obs-err', help='Height levels for observation errors', type=float, nargs='+')
	parser.add_argument('--obs-err-u', type=float, nargs='+')
	parser.add_argument('--obs-err-v', type=float, nargs='+')
	parser.add_argument('-j', '--json', help='JSON configuration file')
	parser.add_argument('-n', '--ntasks', help='Process number', type=int, default=4)
	args = parser.parse_args()

	if args.json:
		config = parse_config(args.json)
		if not args.nature_run_root:
			args.nature_run_root = config['custom']['nature_run']['dir_pattern']
		if not args.nature_run_type:
			args.nature_run_type = config['custom']['nature_run']['type']
		if not has_key(config, ('custom', 'sim_obs', 'profiler', 'obs_err', 'hlev')):
			cli.error('There is no height levels of observation errors in the configure file!')
		loc = np.array(config['custom']['sim_obs']['profiler']['location'])
		args.sid = [str(int(x)).zfill(5) for x in np.linspace(1, len(loc), len(loc))]
		args.lon = loc[:,0]
		args.lat = loc[:,1]
		args.hlev_obs_err = config['custom']['sim_obs']['profiler']['obs_err']['hlev']
		args.obs_err_u = config['custom']['sim_obs']['profiler']['obs_err']['u']
		args.obs_err_v = config['custom']['sim_obs']['profiler']['obs_err']['v']
	if args.hlev_obs_err == None:
		cli.error('Height levels of observation errors should be provided!')

	if args.zlev == None:
		args.zlev = default_zlev
	else:
		args.zlev = np.array(args.zlev) * units.m

	sim_profiler(args.nature_run_root, args.nature_run_type,
		args.sid, args.lon, args.lat, args.zlev,
		args.start_time, args.end_time,
		args.hlev_obs_err, args.obs_err_u, args.obs_err_v, args.ntasks,
		args.output_littler)
